#pragma warning(disable:4661)

#include <map>
#include "IcemVolume.h"
#include "IcemStack.h"
#include <ctime>

#include <fstream>

#include "CudaIcemVolume.h"
#include "cuda.h"
#include "cuda_runtime.h"

using namespace std;


template<class T>
IcemStack<T>::~IcemStack(){
	for (int i=0;i<3;++i)
		cudaFree(cudaImage[i]);
}


//compute slice and return pointer to cuda float3 array on device
template<class T>
void IcemStack<T>::ComputeSlice(double pos[3], int dir){
	short3** imageStack;
	short3** cuImageStack;
	float* alpha;
	float* cuAlpha;
	
	size_t N = volumeStack.size();
	if (talairach)
		++N;
	if (filteredTalairach)
		++N;
	imageStack = (short3**)malloc(N*sizeof(short3*));
	alpha = (float*)malloc(N*sizeof(float));
	cudaMalloc(&cuImageStack,N*sizeof(short3*));
	cudaMalloc(&cuAlpha,N*sizeof(float));
	
	int i=0;
	for (auto it = volumeStack.begin();it != volumeStack.end();++it) {
		alpha[i] = it->second->GetAlpha();
		bool interp = true;
		//if (i==5 ||
		VolumeId id = it->first;
		if (id == mask || id == stat || id == fMRI){
			interp = false;
        }
		imageStack[i] = it->second->GetSlice(pos,dir,minCoordinates,false,interp);
		++i;
	}	
    if (talairach) {
		bool interp = false;
		alpha[i] = talairach->GetAlpha();
		imageStack[i++] = talairach->GetSlice(pos,dir,minCoordinates,false,interp);
	}	
    if (filteredTalairach) {
		bool interp = false;
		alpha[i] = filteredTalairach->GetAlpha();
		imageStack[i++] = filteredTalairach->GetSlice(pos,dir,minCoordinates,false,interp);
	}
	cudaMemcpy(cuImageStack,imageStack,N*sizeof(short3*),cudaMemcpyHostToDevice);
	cudaMemcpy(cuAlpha,alpha,N*sizeof(float),cudaMemcpyHostToDevice);
	int sizeX, sizeY;
	
	GetSliceSize(dir,&sizeX, &sizeY);
	CudaStack<<<NBLOCKS,NTHREADS>>>(N, cudaImage[dir],cuImageStack,cuAlpha, sizeX, sizeY,dir==2);
	int ret = cudaDeviceSynchronize();	
	cudaFree(cuImageStack);
	cudaFree(cuAlpha);
}

template<class T>
void IcemStack<T>::UpdateRenderedVolume() {	size_t pitch = renderedVolume.pitch;
    size_t slicePitch = pitch * size[1];
	double pos[3] = {0,0,0};
	for (size_t i=0;i<size[2];++i) {
		pos[2] = 0;
		ComputeSlice(pos,0);
		int ret;
		ret = cudaMemcpy2D((char*)renderedVolume.ptr + i*slicePitch,pitch,
					  cudaImage[0], size[0]*sizeof(uchar3),
					  size[0]*sizeof(uchar3),size[1], cudaMemcpyDeviceToDevice);
	}
}


//copy slice to host memory
template<class T>
void IcemStack<T>::GetSlice(double pos[3], int dir, uchar3* result){	
	ComputeSlice(pos,dir);
	int sizeX, sizeY;	
	GetSliceSize(dir,&sizeX, &sizeY);
	cudaMemcpy(result,cudaImage[dir], sizeX * sizeY * sizeof(uchar3), cudaMemcpyDeviceToHost );
}


template<class T>
int IcemStack<T>::GetFirstSlice(double pos[3], int dir, uchar3* result){
	short3* cuOut;
	auto it = volumeStack.begin();
	double alphaBackup = it->second->GetAlpha();
	it->second->SetAlpha(1);
	cuOut = it->second->GetSlice(pos,dir,minCoordinates,true,true);
	it->second->SetAlpha(alphaBackup);
	int sizeX, sizeY;
	GetSliceSize(dir, &sizeX, &sizeY);
	short* host_out = (short*)malloc(sizeX*sizeY*sizeof(short3));
	cudaMemcpy(host_out,cuOut,(sizeX) * (sizeY) *sizeof(short3), cudaMemcpyDeviceToHost );
	for (int i=0;i<sizeX*sizeY;++i) {
		result[i].x=(unsigned char)host_out[i*3];
		result[i].y=(unsigned char)host_out[i*3+1];
		result[i].z=(unsigned char)host_out[i*3+2];
	}
	return 0;
}


template<class T>
void IcemStack<T>::SetAlpha(int id, double alpha) {
	VolumeId vid = static_cast<VolumeId>(id);
	if (volumeStack.find(vid) != volumeStack.end()) {
		volumeStack[vid]->SetAlpha(alpha);
	}
	if (vid==atlas)
		talairach->SetAlpha(alpha);
	if (vid == filteredAtlas)
		filteredTalairach->SetAlpha(alpha);
	//UpdateRenderedVolume();
}

template<class T>
void IcemStack<T>::SetContrast(int id, double contrast) {
	VolumeId vid = static_cast<VolumeId>(id);
	if (volumeStack.find(vid) != volumeStack.end()) {
		volumeStack[vid]->SetContrast(contrast);
	}
}

template<class T>
void IcemStack<T>::SetBrightness(int id, double brightness) {
	VolumeId vid = static_cast<VolumeId>(id);
	if (volumeStack.find(vid) != volumeStack.end()) {
		volumeStack[vid]->SetBrightness(brightness);
	}
}

template<class T>
void IcemStack<T>::SetColormap(int id, unsigned char cm[3*256]) {
	VolumeId vid = static_cast<VolumeId>(id);
	if (volumeStack.find(vid) != volumeStack.end()) {
		volumeStack[vid]->SetColormap(cm);
	}
}

template class IcemStack<uint8_t>;